/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: alg remote update.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#include "alg_remote_update.h"
#include "log.h"
#include "dtype_common.h"
#include "adapter_hccp.h"

using namespace std;

namespace hccl {

AlgHdcsRemoteUpdate::AlgHdcsRemoteUpdate() : funcTable_({
    {EmbeddingServiceParam::OperatorType::REMOTE_UPDATE,
        std::bind(&AlgHdcsRemoteUpdate::AlgUpdate, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::WAIT_UPDATE,
        std::bind(&AlgHdcsRemoteUpdate::AlgWaitUpdate, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_UPDATE_REDUCE_SUM,
        std::bind(&AlgHdcsRemoteUpdate::AlgUpdateReduceSum, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_UPDATE_KEY_REDUCE,
        std::bind(&AlgHdcsRemoteUpdate::AlgUpdateKeyReduce, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_UPDATE_SEND_REQUEST,
        std::bind(&AlgHdcsRemoteUpdate::AlgUpdateSendRequest, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_UPDATE_RECV_RESPONSE,
        std::bind(&AlgHdcsRemoteUpdate::AlgUpdateRecvResponse, this, std::placeholders::_1)},
    {EmbeddingServiceParam::OperatorType::REMOTE_UPDATE_RESET_UNIQUE_HANDLE,
        std::bind(&AlgHdcsRemoteUpdate::AlgUpdateResetUniqueHandle, this, std::placeholders::_1)}
    })
{
}

HcclResult AlgHdcsRemoteUpdate::Process(HcclOpPtr opPtr)
{
    CHK_PTR_NULL(opPtr);
    auto op = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    CHK_PRT_RET(funcTable_.find(op->param_.opType) == funcTable_.end(),
        HCCL_ERROR("[Process]funcTable_ without opType[%u]", op->param_.opType), HCCL_E_PARA);
    CHK_RET(funcTable_[op->param_.opType](op));
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgUpdate(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    ReduceSumTask *reduceSumTaskkPtr = new (std::nothrow) ReduceSumTask();
    CHK_PTR_NULL(reduceSumTaskkPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *reduceSumTaskkPtr = param;
    reduceSumTaskkPtr->taskType = OpTaskType::UPDATE;

    OpTaskPtr taskPtr(reduceSumTaskkPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgWaitUpdate(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    u64 offset = 0;
    for (auto &pair : opPtr->param_.embeddingParam.waitLookUpRanks) {
        TransTask *transTaskPtr = new (std::nothrow) TransTask();
        CHK_PTR_NULL(transTaskPtr);
        transTaskPtr->count = pair.second;
        transTaskPtr->dataType = HCCL_DATA_TYPE_INT8; // embedding表数据为fp32
        transTaskPtr->recvMem = static_cast<u8 *>(opPtr->param_.values)
            + offset;
        transTaskPtr->dstRank = pair.first;
        transTaskPtr->taskType = OpTaskType::RECV;
        OpTaskPtr taskPtr(transTaskPtr);
        opPtr->PushTask(std::move(taskPtr));
        offset += pair.second;
    }
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgUpdateReduceSum(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    ReduceSumTask *reduceSumTaskkPtr = new (std::nothrow) ReduceSumTask();
    CHK_PTR_NULL(reduceSumTaskkPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *reduceSumTaskkPtr = param;
    reduceSumTaskkPtr->taskType = OpTaskType::REDUCE_SUM;

    OpTaskPtr taskPtr(reduceSumTaskkPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgUpdateKeyReduce(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    KeyReduceTask *keyReduceTask = new (std::nothrow) struct KeyReduceTask();
    CHK_PTR_NULL(keyReduceTask);
    keyReduceTask->tag = opPtr->param_.tag;
    keyReduceTask->keyMaxNum = opPtr->param_.keyMaxNum;
    keyReduceTask->keys = static_cast<s64 *>(opPtr->param_.keys);
    keyReduceTask->value = opPtr->param_.values;
    keyReduceTask->taskType = OpTaskType::KEY_REDUCE;
    keyReduceTask->valueType = opPtr->param_.valuesType;
    keyReduceTask->valueItemSize = opPtr->param_.embeddingParam.valueItemSize;
    keyReduceTask->keyTransferMem = opPtr->param_.embeddingParam.keyTransferMem;
    keyReduceTask->valueTransferMem = opPtr->param_.embeddingParam.valueTransferMem;

    keyReduceTask->haveRdmaConn = opPtr->param_.haveRdmaConn;
    keyReduceTask->rdmaEnveInfosTransferMem = opPtr->param_.embeddingParam.rdmaEnveInfosTransferMem;

    keyReduceTask->pairedMode = opPtr->param_.pairedMode;
    if (opPtr->param_.pairedMode) {
        keyReduceTask->indices = opPtr->param_.indices;
        keyReduceTask->numUniqued = opPtr->param_.numUniqued;
        keyReduceTask->psSeg = opPtr->param_.psSeg;
        keyReduceTask->psSegNum = opPtr->param_.psSegNum;
    }

    OpTaskPtr taskPtr(keyReduceTask);
    opPtr->PushTask(move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgUpdateSendRequest(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    ReduceSumTask *updateTransTaskPtr = new (std::nothrow) ReduceSumTask();
    CHK_PTR_NULL(updateTransTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *updateTransTaskPtr = param;
    updateTransTaskPtr->taskType = OpTaskType::SEND_REQUEST;

    OpTaskPtr taskPtr(updateTransTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgUpdateRecvResponse(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    ReduceSumTask *updateTransTaskPtr = new (std::nothrow) ReduceSumTask();
    CHK_PTR_NULL(updateTransTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *updateTransTaskPtr = param;
    updateTransTaskPtr->taskType = OpTaskType::RECV_RESPONSE;

    OpTaskPtr taskPtr(updateTransTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}

HcclResult AlgHdcsRemoteUpdate::AlgUpdateResetUniqueHandle(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    ReduceSumTask *updateTransTaskPtr = new (std::nothrow) ReduceSumTask();
    CHK_PTR_NULL(updateTransTaskPtr);
    EmbeddingServiceParam &param = opPtr->param_;
    *updateTransTaskPtr = param;
    updateTransTaskPtr->taskType = OpTaskType::RESET_UNIQUE_HANDLE;

    OpTaskPtr taskPtr(updateTransTaskPtr);
    opPtr->PushTask(std::move(taskPtr));
    return HCCL_SUCCESS;
}
}
